import sys
import collections
from matplotlib.patches import Rectangle, Arrow, FancyArrow
try:
    from matplotlib.backends.backend_cairo import FigureCanvasCairo as FigureCanvas
except:
    from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas

from matplotlib.figure import Figure
from cStringIO import StringIO

import sys
import collections

SCALE = 50

default_colors = (
           ('default', (.26, '#aaaaaa')) # default catches any that's not otherwise defined it's first so it appears at the 'back' of the image.
          ,('gene', (.25, '#ff22aa'))
          ,('MRNA', (.24, '#0000ff'))
          ,('CDS' , (.23, '#008800'))
        )

# http://matplotlib.sourceforge.net/matplotlib.patches.html#Rectangle
class Feature(Rectangle):
    def __init__(self, feature, height=5, color='#aaaaaa', **kwargs):
        strand = feature.get('strand', -height/2.)
        b = 0.2 * strand * SCALE
        b -= height/2. #. * SCALE
        w = feature['stop'] - feature['start']
        #print >>sys.stderr, 'bottom, left, w, h', b, feature['start'], w, height
        Rectangle.__init__(self, (feature['start'], b)
                , width=w, height=height, fill=False, ec=color, **kwargs)


class FArrow(FancyArrow):
    def __init__(self, feature, height=5, ec='#aaaaaa', **kwargs):
        strand = feature.get('strand', -height/2.)
        b = 0.2 * strand * SCALE
        #else: b -= height/2.
        #b -= height/2. #. * SCALE
        #print >>sys.stderr, 'bottom, left, w, h', b, feature['start'], w, height
        x, dx = feature.get('strand') == -1 and (feature['stop'], feature['start'] - feature['stop']) \
                                            or (feature['start'], feature['stop']  - feature['start'])
        y, dy = b, 0
        #Arrow.__init__(self, x, y, dx, dy, width=height * 3, ec=color, fc=color, **kwargs)
        kwargs.setdefault('width', 16)
        kwargs.setdefault('head_width', kwargs['width'] + 4)
        kwargs.setdefault('fc', ec)
        
        kwargs.setdefault('head_length', min(abs(dx), 800))
        FancyArrow.__init__(self, x, y, dx, dy, length_includes_head=True
                , head_starts_at_zero=False, 
                 ec=ec, **kwargs)

from matplotlib.ticker import Formatter
class KFormatter(Formatter):
    def __init__(self, fmt, div):
        self.fmt = fmt
        self.div = div

    def __call__(self, x, pos=None):
        return self.fmt % ( x / self.div)

class Panel(Figure):
    """
    convienence wrapper for drawing a tile.

    >>> t = Panel(dpi=256)
    >>> t.plot([10,20,30],[20,50,20]) #doctest: +ELLIPSIS
    [<matplotlib.lines.Line2D object ...]
    >>> t.plot([[10,20],[3,20],[50,20]]) #doctest: +ELLIPSIS
    [<matplotlib.lines.Line2D object ...]
    >>> t.save('t2.png')
    """
    def __init__(self, bounds=None, dpi=128, width=512
                            , height=256, alpha=0, **kwargs):
 
        Figure.__init__(self, dpi=dpi, frameon=False, **kwargs)
        self.set_size_inches(float(width)/dpi, float(height)/dpi)

        self.figurePatch.set_alpha(alpha)
        self.figurePatch.set_linewidth(0.1)

        self.ax = self.add_axes((0, 0.1, 1, 1), alpha=alpha, frameon=False, yticks=())
        self.set_formatter(bounds)
        self.ax.set_autoscale_on(False)

        self.canvas = FigureCanvas(self)
        if not bounds is None:
            self.ax.set_xlim(bounds['start'], bounds['stop'])
        self.ax.set_ylim(-SCALE, SCALE)
        self.ax.set_yticks([])

    def plot(self, *args, **kwargs):
        return self.ax.plot(*args, **kwargs)

    def set_formatter(self, bounds):
        if bounds is None: return

        rng = bounds['stop'] - bounds['start']
        if rng > 20000 and rng < 1000000:
            self.ax.xaxis.set_major_formatter(KFormatter("%iK", 10000))
        elif rng > 1000000:
            self.ax.xaxis.set_major_formatter(KFormatter("%iM", 1000000))

    def save(self, path):
        return self.canvas.print_figure(path, dpi=self.dpi)


def gplot(features, bounds=None, width=768, height=256, filename=None, colors=None, **kwargs):
    """Plot a list of features or a gendex object.
    
    features: a list of features dict, or a gendex object.
    bounds  : a feature like: {'start': 0, 'stop': 10000} indicating
              the extent of the region to plot. This argument is
              ignored unless features is a gendex object.
    filename: path of a file in which to store the image.
              if filename is sys.stdout ... that will work
              as will a StringIO() object.
    colors  : a list of tuples (not a dict as the order is important)
              of the form: ('feature_type', (height, color))
              e.g:         ('CDS',          (12,  '#ff0000'))
              these indicate both the order to draw the features --
              for instance you may want the CDS to appear above the
              actual gene -- and the height and color they should be
              drawn. see gendex.plot.default_colors for
              an example. these _are_ case sensitive.
              any kwargs are sent in to the matplotlib Rectangle that
              is used for each feature.
    """
    p = Panel(bounds=bounds, width=width, height=height)

    if colors is None: colors = default_colors
    if features.__class__.__name__ in ('Genedex', 'Pairdex'):
        if not bounds is None:     
            features = features.intersection(bounds)
        else:
            features = features._store
    # save the order.
    feat_order = [gtype for gtype, clr in colors]
    colors = dict([(t, (size * SCALE, clr)) for t, (size, clr) in colors])
    rects = collections.defaultdict(list)

    # this still has the problem that sparse areas and dense areas
    # will be rendered differently.
    feats_per_px = 1.0 * len(features)/width
    print >>sys.stderr, feats_per_px
    for f in features:

        fstart, fstop, fstrand = f['start'], f['stop'], f['strand']
        # first look for CDS, exons, etc.
        if feats_per_px < 0.02:
            for ftype in feat_order:
                if not ftype in f: continue
                fheight, fcolor = colors[ftype]
                # f[ftype] looks like:
                # f['CDS'] == [(start, stop), (start, stop...]
                for start, stop in f[ftype]: 
                    # if this cds is near the end of the gene, taper it as an arrow
                    if (fstrand == -1 and start - fstart < 10)  \
                            or (fstrand == 1 and fstop - stop < 10):
                        feat = FArrow({'start': start, 'stop': stop, 'strand': fstrand}
                            , width=fheight, ec=fcolor, head_width=fheight # doesnt have wings.
                            , head_length=min(400, stop-start) # how far to taper
                            , fc=None, **kwargs)

                        rects[ftype].append(feat)
                    else:
                        feat = Feature({'start':start, 'stop':stop, 'strand': fstrand}, fheight, fcolor, **kwargs)
                        rects[ftype].append(feat)

        # always draw the 'start' 'stop' as an arrow
        fheight, fcolor = colors['default']
        feat = FArrow(f, width=fheight, ec=fcolor, **kwargs)
        rects['default'].append(feat)

    for ftype in feat_order:
        for feat in rects[ftype]:
            p.ax.add_patch(feat)
    # and chop to the requested bounds:
    if bounds is None:
        # blech. how to make mpl do this automatically?
        import operator
        features.sort(key=operator.itemgetter('start'))
        p.ax.set_xlim(features[0]['start'], features[-1]['stop'])

    # TODO: return if its a StringIO()
    p.save(filename)

